///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/ham/freeze_core.cpp
///
/// \brief   implementation file for freezing the core
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<string>
#include<vector>
#include<complex>

#include<formic/ham/freeze_core.h>
#include<formic/ham/ham.h>
#include<formic/ham/tei.h>
#include<formic/input/base.h>
#include<formic/exception.h>
#include<formic/output/output.h>
#include<formic/numeric/numeric.h>
#include<formic/mpi/interface.h>
#include<formic/all_data.h>
#include<formic/archive.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   writes the hamiltonian integrals that result from freezing the specified orbitals
///          to a text file
///
/// \param[in]      userinp        user input options
/// \param[in]      ham            The hamiltonian, which holds the integrals.
/// \param[in]      out_name       the name of the output file to write to
/// \param[in]      frz_inds       indices of orbitals that will be frozen
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::freeze_core_to_text_file(const formic::InputBase & userinp,
                                      formic::Hamiltonian<S> & ham,
                                      const std::string & out_name,
                                      const std::vector<int> & frz_inds) {

  // only the root process does anything
  if ( formic::mpi::rank() != 0 )
    return;

  // get the new and old number of orbitals in the basis as well as how many are being frozen
  const int no = userinp.get<int>("nao");
  const int nf = frz_inds.size();
  const int nn = no - nf;

  // check index sanity
  if ( nf >= no )
    throw formic::Exception("formic::freeze_core_to_text_file cannot freeze %i orbitals when there are only %i in the basis") % nf % no;
  for (std::vector<int>::const_iterator it = frz_inds.begin(); it != frz_inds.end(); it++)
    if ( *it < 0 || *it >= no )
      throw formic::Exception("orbital index %i is out of bounds in formic::freeze_core_to_text_file") % (*it);

  // get a set of the to-be-frozen indices
  const std::set<int> frz_set(frz_inds.begin(), frz_inds.end());

  // get array of non-frozen orbital indices
  std::vector<int> not_frz;
  for (int i = 0; i < no; i++)
    if ( frz_set.count(i) == 0 )
      not_frz.push_back(i);
  assert( not_frz.size() == nn );

  // create a new hamiltonian of the correct size
  formic::Hamiltonian<S> new_ham(nn);

  // compute new nuclear and frozen core energy
  S ecn = formic::unity(S()) * ham.e_cor_nuc();
  for (int r = 0; r < nf; r++) {
    const int rr = frz_inds.at(r);
    ecn += 2.0 * ham.get_oei(rr,rr);
  }
  for (int r = 0; r < nf; r++)
  for (int s = 0; s < nf; s++) {
    const int rr = frz_inds.at(r);
    const int ss = frz_inds.at(s);
    ecn += 2.0 * ham.get_tei(rr,rr,ss,ss) - 1.0 * ham.get_tei(rr,ss,ss,rr);
  }
  new_ham.e_cor_nuc() = formic::real(ecn);

  // get new overlap integrals
  for (int i = 0; i < nn; i++)
  for (int j = 0; j < nn; j++) {
    const int ii = not_frz.at(i);
    const int jj = not_frz.at(j);
    new_ham.get_ovi(i,j) = ham.get_ovi(ii,jj);
  }

  // compute new one electron integrals
  for (int i = 0; i < nn; i++)
  for (int j = 0; j < nn; j++) {
    const int ii = not_frz.at(i);
    const int jj = not_frz.at(j);
    new_ham.get_oei(i,j) = ham.get_oei(ii,jj);
    for (int r = 0; r < nf; r++) {
      const int rr = frz_inds.at(r);
      new_ham.get_oei(i,j) += 2.0 * ham.get_tei(ii,jj,rr,rr) - 1.0 * ham.get_tei(ii,rr,rr,jj);
    }
  }

  // get new hamiltonian's two electron integrals
  for (int i = 0; i < nn; i++)
  for (int j = i; j < nn; j++)
  for (int k = 0; k < nn; k++)
  for (int l = k; l < nn; l++) {
    const int ii = not_frz.at(i);
    const int jj = not_frz.at(j);
    const int kk = not_frz.at(k);
    const int ll = not_frz.at(l);
    new_ham.get_tei(i,j,k,l) = ham.get_tei(ii,jj,kk,ll);
  }

  // write the integrals to the specified file
  new_ham.write_text_ints(userinp, out_name);

}

template void formic::freeze_core_to_text_file(const formic::InputBase & userinp,
                                               formic::Hamiltonian<double> & ham,
                                               const std::string & out_name,
                                               const std::vector<int> & frz_inds);
template void formic::freeze_core_to_text_file(const formic::InputBase & userinp,
                                               formic::Hamiltonian<std::complex<double> > & ham,
                                               const std::string & out_name,
                                               const std::vector<int> & frz_inds);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   Function called for the "Make_FC_Text_Ham" command, which writes to disk a new
///          hamiltonian (in text format) in which specified core orbitals have been frozen.
///
/// \param[in]     arch         archive to read command arguments from
/// \param[in]     data         data on which to execute the command
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void formic::Make_FC_Text_Ham_command_function(formic::Archive & arch, formic::AllData & data) {

  // read a name to give the output text file
  std::string out_name;
  formic::mpi::read_and_bcast(arch, out_name, "Failed to read output file name in formic::Make_FC_Text_Ham_command_function.");

  // read list of orbitals to freeze
  std::vector<int> orb_inds;
  formic::mpi::read_and_bcast(arch, orb_inds, "Failed to read indices of orbitals to freeze in formic::Make_FC_Text_Ham_command_function.");

  // freeze the core and write the result to a text file
  if (data.userinp.get<std::string>("scalar_type") == "real") {
    formic::freeze_core_to_text_file(data.userinp, data.hd, out_name, orb_inds);
  } else if (data.userinp.get<std::string>("scalar_type") == "complex") {
    formic::freeze_core_to_text_file(data.userinp, data.hc, out_name, orb_inds);
  } else
    throw formic::Exception("unknown scalar type \"%s\" when executing command \"Make_FC_Text_Ham\"") % data.userinp.get<std::string>("scalar_type");

}
